import os
from typing import List, Optional, Any, Dict

import pandas
from pandas import ExcelWriter

from oneconvert.datasource.base import BaseDataSource


def _flatten_json(data: Any) -> List[Dict[str, Any]]:
    flattened = []

    if isinstance(data, list):
        for item in data:
            if isinstance(item, (dict, list)):
                flattened.extend(_flatten_json(item))
            else:
                flattened.append({"value": item})
    elif isinstance(data, dict):
        items = {"_".join(k.split(".")): v for k, v in _flatten_dict(data).items()}
        flattened.append(items)
    else:
        flattened.append({"value": data})

    return flattened


def _flatten_dict(d: Dict[str, Any], parent_key: str = "") -> Dict[str, Any]:
    items = []
    for k, v in d.items():
        new_key = f"{parent_key}_{k}" if parent_key else k
        if isinstance(v, dict):
            items.extend(_flatten_dict(v, new_key).items())
        elif isinstance(v, list):
            for i, item in enumerate(v):
                list_key = f"{new_key}[{i}]"
                if isinstance(item, dict):
                    items.extend(_flatten_dict(item, list_key).items())
                else:
                    items.append((list_key, item))
        else:
            items.append((new_key, v))
    return dict(items)


class ExcelDataSource(BaseDataSource):
    def read(self, sheet_name: str | int | list[str] | list[int] | None = 0, header: int = 0, skip_rows=0,
             use_cols: Optional[List[str]] = None, **kwargs):
        """
        读取EXCEL表格数据并转JSON
        :param sheet_name: 读取的sheet页，默认第1页，即sheet_name=0
        :param header:  表头行索引（0 开始）
        :param skip_rows: 跳过前 N 行（如注释）
        :param use_cols: 读取的列
        :param kwargs:
        :return: List | Dict
                 查询单个sheet，返回数据对象列表
                 查询多个sheet，返回字典，key为sheet名称，value为数据对象列表
        """
        if not os.path.isfile(self.uri):
            raise ValueError(f'the file {self.uri} not found')

        df = pandas.read_excel(self.uri,
                               sheet_name=sheet_name,
                               header=header,
                               skiprows=skip_rows,
                               usecols=use_cols)

        if isinstance(df, dict):
            result = {}
            for key, value in df.items():
                result[key] = list(value.to_dict(orient="index").values())
        else:
            result = list(df.to_dict(orient="index").values())
        return result

    def write(self, data, sheet_name: str = "Sheet1", flatten: bool = False,
              index: bool = False, **kwargs):
        """
        转化JSON数据为excel表格，并保存
        :param data: JSON数据
        :param sheet_name: 工作表名称，单个工作表，即data是列表时生效
        :param flatten: 是否扁平化处理数据。自动将嵌套结构展开（如 {"a": {"b": 1}} 转为 a_b: 1）
        :param index: 是否保留索引
        :param kwargs:
        :return:
        """
        dir_name = os.path.dirname(self.uri)
        if dir_name and not os.path.exists(dir_name):
            os.makedirs(dir_name, exist_ok=True)

        with ExcelWriter(self.uri, engine="openpyxl") as writer:
            if isinstance(data, dict) and all(isinstance(v, list) for v in data.values()):
                for sheet, sheet_data in data.items():
                    sheet_data = _flatten_json(sheet_data) if flatten else sheet_data
                    df = pandas.DataFrame(sheet_data)
                    df.to_excel(writer, sheet_name=sheet, index=index)
            else:
                sheet_data = _flatten_json(data) if flatten else data
                df = pandas.DataFrame(sheet_data)
                df.to_excel(writer, sheet_name=sheet_name, index=index)
